# 实现Test1
from dolfinx import mesh, plot, fem, default_real_type
from dolfinx.fem import functionspace, petsc
from mpi4py import MPI
import pyvista as pv
import numpy as np
import basix
from basix.ufl import element, mixed_element
import ufl
from petsc4py import PETSc

import tools  # 定义了一些坐标计算相关的函数
# 算例

# 固定压力点


def marker(x): return np.isclose(x[0], 1) & np.isclose(x[1], 1)


def true_velocity_function(x):
    """
    速度真解 函数
    """
    u1 = -4 * x[1] * (4 - x[0]**2)**2 * (4 - x[1]**2)
    u2 = 4 * x[0] * (4 - x[0]**2) * (4 - x[1]**2)**2
    return [u1, u2]


def true_velocity_function_ufl(x):
    """
    速度真解 ufl函数
    """
    return ufl.as_vector(true_velocity_function(x))


def true_pressure_function(x):
    """
    压强真解 函数
    """
    return 150 * np.sin(x[0])


def true_pressure_function_ufl(x):
    """
    压强真解 函数
    """
    return 150 * ufl.sin(x[0])


def lambda_function(x):
    """拉格朗日乘子真解

    Args:
        x (_type_): msh
    """
    return [np.exp(x[0]), np.exp(x[1])]


def lambda_function_ufl(x):
    """拉格朗日乘子真解
        ufl形式
    Args:
        x (_type_): msh
    """
    return ufl.as_vector([ufl.exp(x[0]), ufl.exp(x[1])])


def laplace_u(x):
    """拉普拉斯算子作用于速度真解
    """
    u1 = 8 * x[1] * (3 * x[0]**4 + 6 * x[0]**2 * x[1]
                     ** 2 - 48*x[0]**2 - 8*x[1]**2 + 80)
    u2 = -8 * x[0] * (6 * x[0]**2 * x[1]**2 - 8 * x[0] **
                      2 + 3 * x[1]**4 - 48 * x[1]**2 + 80)
    return ufl.as_vector([u1, u2])


def grad_p(x):
    """压力梯度真解
    """
    return [150 * np.cos(x[0]), 0]  # 压强真解的梯度，y方向为0


def grad_p_ufl(x):
    """压力梯度真解
    """
    return ufl.as_vector([150 * ufl.cos(x[0]), 0])  # 压强真解的梯度，y方向为0


def displacement_function_ufl(x):
    """位移真解
    """
    return true_velocity_function_ufl(x)  # 位移真解和速度真解相同


def displacement_function(x):
    """位移真解
    """
    return true_velocity_function(x)  # 位移真解和速度真解相同


def c(mu, Y):
    """双线性形式 c(μ, Y) 这里ufl.sym()的使用存在困惑 是否有必要?
    """
    A = (ufl.inner(ufl.grad(mu), ufl.grad(Y))) * dx_solid
    B = ufl.inner(mu, Y) * dx_solid
    return A
    # return (ufl.inner(ufl.sym(ufl.grad(mu)), ufl.sym(ufl.grad(Y))) + ufl.inner(mu, Y)) * dx_solid


def d(x):
    # 位移约束 d = u(X̄) - X
    A = np.array(true_velocity_function(
        identity_map(displacement_function(x))))
    B = np.array(displacement_function(x))
    return A - B


def g(x):
    """固体域体力项 g = -ΔX - λ + g_∂B UFL表达式"""
    neg_laplace_X = -laplace_u(x)  # UFL支持直接对向量取负
    neg_lambda = -lambda_function_ufl(x)

    # 计算 g_∂B = ∇X · n_solid_vec
    X_expr = displacement_function(x)
    grad_X = ufl.grad(displacement_function_ufl(x))
    # g_boundary = ufl.dot(grad_X, n_solid_vec)

    # 组合各项（UFL向量运算）
    # return neg_laplace_X + neg_lambda + g_boundary
    return neg_laplace_X + neg_lambda


def identity_map(s):
    return s  # s 为固体参考域坐标 (x, y)


def fluid_boundary(x):
    """流体边界"""
    return np.isclose(x[0], 2.0) | np.isclose(
        x[0], -2.0) | np.isclose(x[1], -2.0) | np.isclose(x[1], 2.0)


def solid_boundary(x):
    """固体边界"""
    return np.isclose(x[0], 1.0) | np.isclose(
        x[0], -1.0) | np.isclose(x[1], -1.0) | np.isclose(x[1], 1.0)


def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    # 需要导入
    # from petsc4py import PETSc
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data


h_solid = 0.125  # 固体网格大小
h_fluid = 0.125
Lx_solid, Ly_solid = 2, 2  # 固体网格长宽
Lx_fluid, Ly_fluid = 4, 4  # 流体网格长宽
n_solid = [int(Lx_solid / h_solid), int(Ly_solid / h_solid)]
n_fluid = [int(Lx_fluid / h_fluid), int(Ly_fluid / h_fluid)]

solid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-1, -1], [1, 1]], n=n_solid, cell_type=mesh.CellType.triangle)
fluid_msh = mesh.create_rectangle(
    MPI.COMM_WORLD, [[-2, -2], [2, 2]], n=n_fluid, cell_type=mesh.CellType.triangle)
n_solid_vec = ufl.FacetNormal(solid_msh)  # 固体域 单位外法向量


# 定义有限元空间 https://defelement.org/elements/p1-iso-p2.html
u_p2 = element(
    "Lagrange", fluid_msh.basix_cell(), degree=2, shape=(fluid_msh.geometry.dim,), dtype=default_real_type
)
p_p1 = element("Lagrange", fluid_msh.basix_cell(),
               degree=1, dtype=default_real_type)
X_p1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
lambda_p1 = element(
    "Lagrange", solid_msh.basix_cell(), degree=1, shape=(solid_msh.geometry.dim,), dtype=default_real_type
)
dx_fluid = ufl.Measure("dx", domain=fluid_msh)
dx_solid = ufl.Measure("dx", domain=solid_msh)
V = functionspace(fluid_msh, u_p2)  # 流体速度函数空间
Q = functionspace(fluid_msh, p_p1)  # 流体压力函数空间
W = functionspace(solid_msh, X_p1)  # 固体位移函数空间
S = functionspace(solid_msh, lambda_p1)  # 固体拉格朗日乘子函数空间

# 流体域测试函数和试函数
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(X, lambda_) = ufl.TrialFunction(W), ufl.TrialFunction(S)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)
(Y, mu) = ufl.TestFunction(W), ufl.TestFunction(S)

u_bc = fem.Function(V)  # 真实速度函数空间
p_bc = fem.Function(Q)  # 真实压力函数空间
u_bc.interpolate(true_velocity_function)
p_bc.interpolate(true_pressure_function)

# 设置流体边界自由度
fdim = fluid_msh.topology.dim - 1
facets = mesh.locate_entities_boundary(fluid_msh, fdim, fluid_boundary)
v_dofs = fem.locate_dofs_topological(V, fdim, facets)
p_dofs = fem.locate_dofs_geometrical(Q, marker)
bc_vel = fem.dirichletbc(u_bc, v_dofs)  # 速度边界条件
bc_pressure = fem.dirichletbc(p_bc, p_dofs)  # 压力边界条件
bcs = [bc_vel, bc_pressure]  # 边界条件列表

x_fluid = ufl.SpatialCoordinate(fluid_msh)  # 流体域上的空间坐标
x_solid = ufl.SpatialCoordinate(solid_msh)  # 固体域上的空间坐标

# 流体域体力项 f（包含 -Δu + ∇p + c(λ, v)）
# 这里需要实现第五章的算法
# f_fluid = -ufl.div(ufl.grad(u)) + ufl.grad(p) + c(lambda_, v)

# 固体域体力项 g（包含 -ΔX - λ + g_∂B）

# g_boundary = ufl.dot(ufl.grad(X), n_solid_vec)
# g_solid = -ufl.div(ufl.grad(X)) - lambda_ + g_boundary  # g_boundary 需提前定义


# 双线性形 这里ufl.sym()的使用存在困惑 是否有必要?
# 4 * 4矩阵 a_{i}_{j} i, j = 1, 2, 3, 4
a_1_1 = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx_fluid  # a_f(u, v) a_1_1
a_1_2 = -ufl.inner(ufl.div(v), p) * dx_fluid  # (p, v)项
a_1_3 = None
a_1_4 = None  # 暂时没实现

a_2_1 = -ufl.inner(ufl.div(u), q) * dx_fluid
a_2_2 = fem.Constant(fluid_msh, 0.0) * ufl.inner(p, q) * dx_fluid  # (p, q)项
a_2_3 = None
a_2_4 = None

a_3_1 = None
a_3_2 = None
a_3_3 = ufl.inner(ufl.sym(ufl.grad(X)), ufl.sym(
    ufl.grad(Y))) * dx_solid  # a_s(X, Y)
a_3_4 = -c(lambda_, Y)  # c(λ, Y)

a_4_1 = None  # 暂时没实现
a_4_2 = None
a_4_3 = -c(mu, X)  # c(μ, X)
a_4_4 = fem.Constant(solid_msh, 0.0) * ufl.inner(lambda_, mu) * dx_solid

# c(mu, Y)形式 使用函数c(mu, Y)来表示双线性形式
# c（mu, v)形式 Chapter5 给出计算方法
# c1 = c(lambda_, v)  # c(λ, v)
# c2 = -c(lambda_, Y) # c(λ, Y)
# c3 = c(mu, d) # c(μ, u - X)

a = fem.form(
    [
        [a_1_1, a_1_2, a_1_3, a_1_4],
        [a_2_1, a_2_2, a_2_3, a_2_4],
        [a_3_1, a_3_2, a_3_3, a_3_4],
        [a_4_1, a_4_2, a_4_3, a_4_4],
    ]
)
# L = [f 0 g d]^T
# 先进行L1占位
L1 = fem.Constant(fluid_msh, 0.0) * v[0] * dx_fluid  # x方向体力项
L1 += fem.Constant(fluid_msh, 0.0) * v[1] * dx_fluid  # y方向体力项


# L1 = ufl.inner(-laplace_u(x_fluid), v) * dx_fluid
# L1 += ufl.inner(grad_p_ufl(x_fluid), v) * dx_fluid
# L1 += c(lambda_function_ufl(x_fluid), v)  # c(λ, v)
L2 = ufl.inner(fem.Constant(fluid_msh, 0.0), q) * dx_fluid
L3 = ufl.inner(g(x_solid), Y) * dx_solid  # 固体域体力项 省略了g_partialB
# case1 是特殊的 固体网格嵌套在流体网格中 d=0
L4 = fem.Constant(solid_msh, 0.0) * mu[0] * dx_solid  # x方向体力项
L4 += fem.Constant(solid_msh, 0.0) * mu[1] * dx_solid  # y方向体力项


L = fem.form([L1, L2, L3, L4])  # 载荷向量
# 跨网格的积分怎么处理


A = petsc.assemble_matrix_block(a, bcs=bcs)  # petsc4py.PETSc.Mat
L = petsc.assemble_vector_block(L, a, bcs=bcs)  # petsc4py.PETSc.Vec
A.assemble()
L.assemble()
load_vec = L.getArray()


# Algorithm 5.2 无交集算法
fluid_msh.topology.create_connectivity(2, 0)  # 创建流体域顶点和边的连接
solid_msh.topology.create_connectivity(2, 0)  # 创建固体域顶点和边的连接
lambda_dofmap = S.dofmap  # 拉格朗日乘子的自由度映射 因为计算的是c(lambda, v)
v_dofmap = V.dofmap  # 流体速度的自由度映射

num_solid_cells = 2 * n_solid[0]**2  # 2*N_solid*N_solid
num_fluid_cells = 2 * n_fluid[0]**2  # 2*N_fluid*N_fluid

# 计算基函数
gauss_pt1 = [0.6, 0.2, 0.2]  # 标准单元的高斯积分点
gauss_pt2 = [0.2, 0.6, 0.2]
gauss_pt3 = [0.2, 0.2, 0.6]
gauss_pt = np.array(
    [[0.6, 0.2, 0.2],
     [0.2, 0.6, 0.2],
     [0.2, 0.2, 0.6],
     ]
)
# 参考./index-rw-demo/topo.ipynb文件对矢量元tabulate函数的分析
basis_values = lambda_p1.tabulate(0, np.array(gauss_pt[:, :-1]))

# basis_values = lambda_p1.tabulate(0, gauss_pt[:, :-1])
for i in range(num_solid_cells):
    global_solid_cell_idx = lambda_dofmap.cell_dofs(
        i)  # 获取固体单元的全局自由度索引 P1元 长度为3
    solid_cell_points = S.tabulate_dof_coordinates(
    )[global_solid_cell_idx]  # 获取固体单元的自由度坐标 P1元 对应顶点
    T_s_area = tools.get_triangle_area(*solid_cell_points[:, :-1])  # 获取固体单元的面积 *号表示解包


print("f")
# P2/P1/P1/P1 流体32*32 固体16*16剖分 刚度矩阵 10695*10695
# res = getStiffnessMatrix(A)

# # 流体域线性形式
# L_f = ufl.inner(f_fluid, v) * dx_fluid

# # 固体域线性形式
# L_s = ufl.inner(g_solid, Y) * dx_solid + ufl.inner(d, mu) * dx_solid


# # 全局双线性形式
# a = a_f + a_s + c1 + c2 + c3
# a -= ufl.div(v) * p * dx_fluid  # 流体域的压力项
# a -= ufl.div(u) * q * dx_fluid  # 质量守恒
# # 全局线性形式
# L = L_f + L_s

# # 组装矩阵和向量

# b = fem.petsc.assemble_vector(L)
# fem.petsc.apply_lifting(b, [a], [bcs])
# b.ghostUpdate(addv=fem.petsc.ScatterMode.ADD,
#               mode=fem.petsc.ScatterMode.REVERSE)
# fem.petsc.set_bc(b, bcs)

# # 创建并配置Krylov子空间求解器
# solver = fem.petsc.LinearProblem(
#     a, L, bcs=bcs,
#     petsc_options={"ksp_type": "gmres", "pc_type": "lu"}
# )

# # 求解问题
# solution = solver.solve()
